from torch.utils.data import Dataset
from csv import reader as CSV_reader
from os.path import join as o_p_join
from PIL import Image
from torchvision.transforms import\
    Compose, Resize, InterpolationMode, CenterCrop, ToTensor, Normalize


PAIR_DIR = 'Protocol'
PAIR_FILE = 'Pair_list_F.txt'
LIST_FILE = 'Data/list_name.txt'
NUM_CLASSES = 500
DL_BSIZE = 512


class CFP_Dset(Dataset):
    def __init__(self, root_dir, train=True):
        super(CFP_Dset, self).__init__()
        # 创建类型列表
        self.classes = []
        with open(o_p_join(root_dir, LIST_FILE)) as csv_f:
            csv_r = CSV_reader(csv_f)
            for __ in range(NUM_CLASSES):
                row_data = next(csv_r)
                self.classes.append(row_data[0])
        # 统计各项数据列表
        self.root_dir = o_p_join(root_dir, PAIR_DIR)
        self.item_dir = []
        self.item_class = []
        self.item_count = 0
        with open(o_p_join(self.root_dir, PAIR_FILE)) as csv_f:
            csv_r = CSV_reader(csv_f, delimiter=' ')
            for i in range(NUM_CLASSES):
                for j in range(10):
                    row_data = next(csv_r)
                    if (train and (j == 9)) or ((not train) and (j != 9)):
                        continue
                    self.item_count += 1
                    self.item_dir.append(row_data[1])
                    self.item_class.append(i)
        # 确定数据预处理方式
        self.item_trans = Compose([Resize(227, interpolation=InterpolationMode.BICUBIC),
                                CenterCrop(227),
                                ToTensor(),
                                Normalize([0.5, 0.5, 0.5], [0.2, 0.2, 0.2])])

    def __getitem__(self, idx):
        img_dir = o_p_join(self.root_dir, self.item_dir[idx])
        img = self.item_trans(Image.open(img_dir))
        return img, self.item_class[idx]

    def __len__(self):
        return self.item_count
